"""This module provides example tools for web scraping and search functionality.

It includes a basic Tavily search function (as an example)

These tools are intended as free examples to get started. For production use,
consider implementing more robust and specialized tools tailored to your needs.
"""

from typing import Any, Callable, List, Optional, cast
import os

from langchain_community.tools.tavily_search import TavilySearchResults
from langchain_core.runnables import RunnableConfig
from langchain_core.tools import InjectedToolArg
from typing_extensions import Annotated

from react_agent.configuration import Configuration
from react_agent.utils import load_text_db, append_text_db, load_img_db, append_img_db, image_uri_extract


async def web_search(
    query: str, *, config: Annotated[RunnableConfig, InjectedToolArg]
) -> Optional[list[dict[str, Any]]]:
    """Search for general web results.

    This function performs a search using the Tavily search engine, which is designed
    to provide comprehensive, accurate, and trusted results. It's particularly useful
    for answering questions about current events.
    """
    configuration = Configuration.from_runnable_config(config)
    wrapped = TavilySearchResults(max_results=configuration.max_search_results)
    result = await wrapped.ainvoke({"query": query})
    return cast(list[dict[str, Any]], result)


async def text_retrieve(
    query: str, *, config: Annotated[RunnableConfig, InjectedToolArg]
) -> str:
    """从敦煌文本知识库里面检索莫高窟等相关知识，作为智能问答的上下文。

    该函数在知识库中执行文本检索：文本进，文本出。旨在提供全面、准确和可信赖的结果。
    它特别适用于回答关于敦煌莫高窟的问题。
    非敦煌莫高窟的问题，不要调用 text_retrieve !
    """
    configuration = Configuration.from_runnable_config(config)
    db_path = configuration.text_db_path
    vector_store = None
    if os.path.exists(db_path) and os.path.isdir(db_path) and os.listdir(db_path):
        vector_store = load_text_db(db_path)
        print("Load Text db done.")
    else:
        print("First time loading text db. This may take a while.")
        text_dir = db_path.replace("dunhuang_text_db", "dunhuang_raw")
        vector_store = append_text_db(db_path, text_dir)

    docs = vector_store.similarity_search(query, k=configuration.text_top_k)
    docs_content = "上下文：" + "\n\n".join(doc.page_content for doc in docs)
    return docs_content


async def image_retrieve(
    img_uri: str, *, config: Annotated[RunnableConfig, InjectedToolArg]
) -> str:
    # Passed img_uri by llm is Wrong! mismatch original path
    """从敦煌图片知识库里面检索莫高窟等相关图片内容描述，作为智能问答的上下文。

    该函数在知识库中执行图片检索：图片URI(PATH)进，图片内容文本出。旨在提供全面、准确和可信赖的结果。
    它特别适用于回答关于敦煌莫高窟壁画图片的问题。
    非敦煌莫高窟图片的问题，不要调用 image_retrieve !

    Args:
        img_uri (str): Image URI. You need extract the correct path or url value from Question to pass it. Don't change anything in original value!
            For example, Question - "img - './data/dunhuang_raw/0254/莫高窟第254窟  主室 南壁.png'", the correct extracted img_uri is "./data/dunhuang_raw/0254/莫高窟第254窟  主室 南壁.png".
        config (Annotated[RunnableConfig, InjectedToolArg]): The configuration for the agent.

    Returns:
        str: Image content description
    """
    configuration = Configuration.from_runnable_config(config)
    db_path = configuration.img_db_path
    vector_store = None
    if os.path.exists(db_path) and os.path.isdir(db_path) and os.listdir(db_path):
        vector_store = load_img_db(db_path)
        print("Load Image db done.")
    else:
        print("First time loading text db. This may take a while.")
        img_dir = db_path.replace("dunhuang_img_db", "dunhuang_raw")
        vector_store = append_img_db(db_path, img_dir)

    if img_uri.startswith("./data"):
        img_uri = os.path.join(configuration.museumai_root,
                               img_uri.replace("./", "")).replace("\\", "/")
    if not os.path.exists(img_uri):
        raise FileNotFoundError(f"Image {img_uri} not found. Please re-input one.")
    if not os.environ.get("NOMIC_API_KEY"):
        raise ValueError(
            "NOMIC_API_KEY not found and we need it to embed image. Please set it in the environment variable or in .env file")
    docs = vector_store.similarity_search_by_image(img_uri, k=configuration.img_top_k)
    docs_content = []
    for doc in docs:
        img_path = doc.metadata["path"]
        print("Similar image path:", img_path)
        txt_path = img_path.replace(".png", ".txt")
        print("Context text path:", txt_path)
        txt_path = configuration.museumai_root + txt_path.replace("./", "/")
        with open(txt_path, "r", encoding='utf-8') as f:
            img_content = f.read()
            docs_content.append(img_content)
    docs_content = "上下文：" + "\n\n".join(docs_content)
    return docs_content


async def image_retrieve_new(
    query: str, config: Annotated[RunnableConfig, InjectedToolArg]
) -> str:
    # Passed query by llm is Wrong too! No path included
    """从敦煌图片知识库里面检索莫高窟等相关图片内容描述，作为智能问答的上下文。

    该函数在知识库中执行图片检索：接受用户query，抽取其中图片URI并进行图片检索，最后返回图片内容文本。旨在提供全面、准确和可信赖的结果。
    它特别适用于回答关于敦煌莫高窟壁画图片的问题。
    非敦煌莫高窟图片的问题，不要调用 image_retrieve_new !
    """
    img_uri = image_uri_extract(query)
    if img_uri is None:
        raise ValueError(f"No image URI found in the query '{query}'")
    print("Extracted Image URI:", img_uri)
    if img_uri.startswith("./data"):
        img_uri = os.path.join(configuration.museumai_root,
                               img_uri.replace("./", "")).replace("\\", "/")
        print(img_uri)
    if not os.path.exists(img_uri):
        raise FileNotFoundError(f"Image {img_uri} not found. Please re-input one.")

    configuration = Configuration.from_runnable_config(config)
    db_path = configuration.img_db_path
    vector_store = None
    if os.path.exists(db_path) and os.path.isdir(db_path) and os.listdir(db_path):
        vector_store = load_img_db(db_path)
        print("Load Image db done.")
    else:
        print("First time loading text db. This may take a while.")
        img_dir = db_path.replace("dunhuang_img_db", "dunhuang_raw")
        vector_store = append_img_db(db_path, img_dir)

    if not os.environ.get("NOMIC_API_KEY"):
        raise ValueError(
            "NOMIC_API_KEY not found and we need it to embed image. Please set it in the environment variable or in .env file")
    docs = vector_store.similarity_search_by_image(img_uri, k=configuration.img_top_k)
    docs_content = []
    for doc in docs:
        img_path = doc.metadata["path"]
        print("Similar image path:", img_path)
        txt_path = img_path.replace(".png", ".txt")
        print("Context text path:", txt_path)
        txt_path = os.path.join(configuration.museumai_root, txt_path.replace("./", ""))
        with open(txt_path, "r", encoding='utf-8') as f:
            img_content = f.read()
            docs_content.append(img_content)

    docs_content = "上下文：" + "\n\n".join(docs_content)
    return docs_content

async def empty_tool() -> str:
    """该函数表示用户意图识别的一种特殊状态。当用户提出一个相对宽泛的问题时，请调用 empty_tool 函数。
    它返回空的上下文。
    """
    return ""

TOOLS: List[Callable[..., Any]] = [text_retrieve, empty_tool]
